//
//  ViewController.swift
//  AIPhotoAnimation
//
//  Created by Cobbly on 2025/2/17.
//

import UIKit
import CoreML
import Vision
import Photos
import SnapKit
import CoreImage
import CoreGraphics

class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavigationControllerDelegate {
    
    private let imageView: UIImageView = {
        let iv = UIImageView()
        iv.contentMode = .scaleAspectFit
        iv.backgroundColor = .systemGray6
        return iv
    }()
    
    private let selectButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("选择图片", for: .normal)
        button.backgroundColor = .systemBlue
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        return button
    }()
    
    private let processButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("动漫化", for: .normal)
        button.backgroundColor = .systemGreen
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        button.isEnabled = false
        return button
    }()
    
    private let denoiseButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("降噪处理", for: .normal)
        button.backgroundColor = .systemOrange
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        button.isEnabled = false
        return button
    }()
    
    // 添加新的按钮属性
    private let hayaoButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("宫崎骏风格", for: .normal)
        button.backgroundColor = .systemPurple
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        button.isEnabled = false
        return button
    }()
    
    private let shinkaiButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("新海诚风格", for: .normal)
        button.backgroundColor = .systemIndigo
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        button.isEnabled = false
        return button
    }()
    
    // 添加复原按钮
    private let resetButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("复原", for: .normal)
        button.backgroundColor = .systemRed
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        button.isEnabled = false
        return button
    }()
    
    // 添加视频处理按钮
    private let videoButton: UIButton = {
        let button = UIButton(type: .system)
        button.setTitle("视频动漫化", for: .normal)
        button.backgroundColor = .systemTeal
        button.setTitleColor(.white, for: .normal)
        button.layer.cornerRadius = 10
        return button
    }()
    
    // 添加模型属性
    private var denoiseModel: anime_noise0_model?
    private var hayaoModel: AnimeGANv3_Hayao_36?
    private var shinkaiModel: AnimeGANv3_Shinkai_37?
    
    // 添加原始图片属性
    private var originalImage: UIImage?
    
    // 添加历史记录属性
    private var processedImages: [String: UIImage] = [:]
    
    // 使用懒加载方式初始化处理器
    private lazy var hayaoProcessor: HayaoStyleProcessor? = {
        do {
            let processor = try HayaoStyleProcessor()
            print("宫崎骏风格处理器初始化成功")
            return processor
        } catch {
            print("宫崎骏风格处理器初始化失败: \(error)")
            return nil
        }
    }()
    
    private lazy var shinkaiProcessor: ShinkaiStyleProcessor? = {
        do {
            let processor = try ShinkaiStyleProcessor()
            print("新海诚风格处理器初始化成功")
            return processor
        } catch {
            print("新海诚风格处理器初始化失败: \(error)")
            return nil
        }
    }()
    
    private lazy var denoiseProcessor: DenoiseProcessor? = {
        do {
            let processor = try DenoiseProcessor()
            print("降噪处理器初始化成功")
            return processor
        } catch {
            print("降噪处理器初始化失败: \(error)")
            return nil
        }
    }()
    
    override func viewDidLoad() {
        super.viewDidLoad()
        setupUI()
    }
    
    private func setupUI() {
        view.backgroundColor = .white
        view.addSubview(imageView)
        view.addSubview(selectButton)
        view.addSubview(processButton)
        view.addSubview(denoiseButton)
        view.addSubview(hayaoButton)
        view.addSubview(shinkaiButton)
        view.addSubview(resetButton)
        view.addSubview(videoButton)
        
        // 使用 SnapKit 设置约束
        imageView.snp.makeConstraints { make in
            make.top.equalTo(view.safeAreaLayoutGuide).offset(20)
            make.left.equalToSuperview().offset(20)
            make.right.equalToSuperview().offset(-20)
            make.height.equalTo(imageView.snp.width)
        }
        
        // 第一行按钮：选择图片、动漫化、降噪处理
        selectButton.snp.makeConstraints { make in
            make.top.equalTo(imageView.snp.bottom).offset(20)
            make.left.equalTo(view).offset(20)
            make.width.equalTo((view.frame.width - 80) / 3) // 三个按钮等宽
            make.height.equalTo(44)
        }
        
        processButton.snp.makeConstraints { make in
            make.top.equalTo(selectButton)
            make.left.equalTo(selectButton.snp.right).offset(20)
            make.width.equalTo(selectButton)
            make.height.equalTo(44)
        }
        
        denoiseButton.snp.makeConstraints { make in
            make.top.equalTo(selectButton)
            make.left.equalTo(processButton.snp.right).offset(20)
            make.width.equalTo(selectButton)
            make.height.equalTo(44)
        }
        
        // 第二行按钮：宫崎骏风格、新海诚风格、复原
        hayaoButton.snp.makeConstraints { make in
            make.top.equalTo(selectButton.snp.bottom).offset(20)
            make.left.equalTo(view).offset(20)
            make.width.equalTo(selectButton)
            make.height.equalTo(44)
        }
        
        shinkaiButton.snp.makeConstraints { make in
            make.top.equalTo(hayaoButton)
            make.left.equalTo(hayaoButton.snp.right).offset(20)
            make.width.equalTo(selectButton)
            make.height.equalTo(44)
        }
        
        resetButton.snp.makeConstraints { make in
            make.top.equalTo(hayaoButton)
            make.left.equalTo(shinkaiButton.snp.right).offset(20)
            make.width.equalTo(selectButton)
            make.height.equalTo(44)
        }
        
        // 添加视频按钮
        videoButton.snp.makeConstraints { make in
            make.top.equalTo(hayaoButton.snp.bottom).offset(20)
            make.left.equalTo(view).offset(20)
            make.width.equalTo((view.frame.width - 80) / 3)
            make.height.equalTo(44)
        }
        
        // 按钮事件保持不变
        selectButton.addTarget(self, action: #selector(selectButtonTapped), for: .touchUpInside)
        processButton.addTarget(self, action: #selector(processButtonTapped), for: .touchUpInside)
        denoiseButton.addTarget(self, action: #selector(denoiseButtonTapped), for: .touchUpInside)
        hayaoButton.addTarget(self, action: #selector(hayaoButtonTapped), for: .touchUpInside)
        shinkaiButton.addTarget(self, action: #selector(shinkaiButtonTapped), for: .touchUpInside)
        resetButton.addTarget(self, action: #selector(resetButtonTapped), for: .touchUpInside)
        videoButton.addTarget(self, action: #selector(videoButtonTapped), for: .touchUpInside)
    }
    
    @objc private func selectButtonTapped() {
        let imagePicker = UIImagePickerController()
        imagePicker.delegate = self
        imagePicker.sourceType = .photoLibrary
        present(imagePicker, animated: true)
    }
    
    @objc private func processButtonTapped() {
        guard let originalImage = self.originalImage else {
            print("没有原始图片")
            return
        }
        
        processImage(using: denoiseProcessor, image: originalImage, button: processButton)
    }
    
    @objc private func denoiseButtonTapped() {
        guard let originalImage = self.originalImage else {
            print("没有原始图片")
            return
        }
        
        processImage(using: denoiseProcessor, image: originalImage, button: denoiseButton)
    }
    
    // UIImagePickerControllerDelegate 方法
    func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
        picker.dismiss(animated: true)
        
        guard let image = info[.originalImage] as? UIImage else {
            print("无法获取选中的图片")
            return
        }
        
        // 保存原始图片
        originalImage = image
        imageView.image = image
        
        // 启用所有处理按钮
        processButton.isEnabled = true
        denoiseButton.isEnabled = true
        hayaoButton.isEnabled = true
        shinkaiButton.isEnabled = true
        resetButton.isEnabled = true
        videoButton.isEnabled = true
        
        // 清除历史记录
        processedImages.removeAll()
    }
    
    func imagePickerControllerDidCancel(_ picker: UIImagePickerController) {
        picker.dismiss(animated: true)
    }
    
    // 添加新的处理方法
    @objc private func hayaoButtonTapped() {
        guard let originalImage = self.originalImage else {
            print("没有原始图片")
            return
        }
        
        processImage(using: hayaoProcessor, image: originalImage, button: hayaoButton)
    }
    
    @objc private func shinkaiButtonTapped() {
        guard let originalImage = self.originalImage else {
            print("没有原始图片")
            return
        }
        
        processImage(using: shinkaiProcessor, image: originalImage, button: shinkaiButton)
    }
    
    private func processImage(using processor: ImageProcessor?, image: UIImage, button: UIButton) {
        guard let processor = processor else {
            print("处理器未初始化")
            return
        }
        
        // 显示加载指示器
        let activityIndicator = UIActivityIndicatorView(style: .large)
        activityIndicator.center = imageView.center
        activityIndicator.startAnimating()
        imageView.addSubview(activityIndicator)
        
        // 禁用按钮
        button.isEnabled = false
        
        // 在后台处理图片
        Task {
            do {
                let processedImage = try await processor.process(image: image)
                
                // 在主线程更新 UI
                await MainActor.run {
                    self.imageView.image = processedImage
                    activityIndicator.removeFromSuperview()
                    button.isEnabled = true
                }
            } catch {
                print("处理失败: \(error)")
                await MainActor.run {
                    activityIndicator.removeFromSuperview()
                    button.isEnabled = true
                }
            }
        }
    }
    
    // 添加复原按钮事件处理
    @objc private func resetButtonTapped() {
        guard let originalImage = self.originalImage else {
            print("没有原始图片")
            return
        }
        
        // 显示加载指示器
        let activityIndicator = UIActivityIndicatorView(style: .large)
        activityIndicator.center = imageView.center
        activityIndicator.startAnimating()
        imageView.addSubview(activityIndicator)
        
        // 在主线程更新 UI
        DispatchQueue.main.async { [weak self] in
            self?.imageView.image = originalImage
            activityIndicator.removeFromSuperview()
            print("已恢复原始图片")
        }
    }
    
    // 在处理完图片后保存结果
    private func saveProcessedImage(_ image: UIImage, forStyle style: String) {
        processedImages[style] = image
    }
    
    // 处理视频按钮点击事件
    @objc private func videoButtonTapped() {
        let videoConverter = VideoAnimeConverter()
        videoConverter.modalPresentationStyle = .fullScreen
        present(videoConverter, animated: true)
    }
}

// UIImage 扩展保持不变...

extension UIImage {
    func toCVPixelBuffer() -> CVPixelBuffer? {
        let width = Int(self.size.width)
        let height = Int(self.size.height)
        
        var pixelBuffer: CVPixelBuffer?
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
                                       width,
                                       height,
                                       kCVPixelFormatType_32BGRA,
                                       nil,
                                       &pixelBuffer)
        
        guard let buffer = pixelBuffer, status == kCVReturnSuccess else {
            return nil
        }
        
        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }
        
        let context = CGContext(data: CVPixelBufferGetBaseAddress(buffer),
                              width: width,
                              height: height,
                              bitsPerComponent: 8,
                              bytesPerRow: CVPixelBufferGetBytesPerRow(buffer),
                              space: CGColorSpaceCreateDeviceRGB(),
                              bitmapInfo: CGImageAlphaInfo.premultipliedFirst.rawValue)
        
        context?.draw(self.cgImage!, in: CGRect(x: 0, y: 0, width: width, height: height))
        
        return buffer
    }
}

